from typing import Optional
from dotenv import load_dotenv
import os
from fastapi.responses import RedirectResponse
from qcloud_cos import CosConfig, CosS3Client
from fastapi import FastAPI, HTTPException, Request, Response
import logging
import random

# 加载.env文件中的环境变量
load_dotenv()
BUCKET = "zq4-1318886512"
COS_SECRET_ID = os.environ["COS_SECRET_ID"]
COS_SECRET_KEY = os.environ["COS_SECRET_KEY"]
COS_REGION = os.environ["COS_REGION"]
EXPIRED = 15

config = CosConfig(
    Region=COS_REGION,
    SecretId=COS_SECRET_ID,
    SecretKey=COS_SECRET_KEY,
    Token=None,
    Scheme="https",
)
client = CosS3Client(config)

app = FastAPI()
register = {}

def getKey():
    return str(random.randint(1000, 9999))
def plain_text_response(content: str):
    return Response(content=content, media_type="text/plain")


@app.get("/generate_presigned_url")
async def generate_presigned_url(object_key: str):
    """
    生成预签名上传 URL
    :param object_key: COS 中的对象键名
    :return: 预签名 URL
    """
    try:
        signed_url = client.get_presigned_url(
            Bucket=BUCKET,
            Key=object_key,
            Expired=EXPIRED,
            Method="PUT",
            Headers={"Access-Control-Allow-Origin": "*"},
        )
        return {"signed_url": signed_url}
    except Exception as e:
        logging.error(f"生成预签名 URL 时发生错误：{e}")
        raise HTTPException(status_code=500, detail="生成预签名 URL 失败")


@app.post("/upload_file")
async def upload_file(request: Request):
    body = await request.body()
    object_key = body.decode()
    key = getKey()
    register[key] = {
        "ty": "file",
        "content": object_key,
    }
    print("key:", key, "object:", object_key)
    return plain_text_response(key)


@app.post("/upload_text")
async def upload_text(request: Request):
    body = await request.body()
    text = body.decode()
    key = getKey()
    register[key] = {
        "ty": "text",
        "content": text,
    }
    return plain_text_response(key)

"""
mode: common, plain, instant
"""
@app.get("/download")
async def download(query: str, mode: str = "common"):
    if query not in register:
        raise HTTPException(status_code=404, detail="Key not found")
    content = register[query]

    if content["ty"] == "file":
        signed_url = client.get_presigned_download_url(
            Bucket=BUCKET, Key=content["content"], Expired=EXPIRED
        )
        if mode == "plain":
            return plain_text_response(signed_url)
        if mode == "instant":
            return RedirectResponse(url=signed_url)
        return {"ty": "file", "signed_url": signed_url, "filename": content["content"]}
    elif content["ty"] == "text":
        if mode == "plain" or mode == "instant":
            return plain_text_response(content["content"])
        return {"ty": "text", "content": content["content"]}
    else:
        raise HTTPException(status_code=500, detail="Unknown type")


@app.middleware("http")
async def CORS(request: Request, call_next):
    response = await call_next(request)
    response.headers["Access-Control-Allow-Origin"] = "*"
    return response


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="127.0.0.1", port=18922)
